When running HuggingFace BERT (any size) fine-tuning tutorial with transformers version >= 4.21.0 and using XLA_USE_BF16=1 or XLA_DOWNCAST_BF16=1, I see NaNs in the loss after the first step.

What does this PR do?

This PR addresses the issue where the model code passes a value that is out of range for XLA_USE_BF16=1 or XLA_DOWNCAST_BF16=1, so the conversion would cast it to -inf.

The NaNs likely come from the transformers library change: #17306 . This PR replaced many lines which used to be -float(inf) (or other small constants) with torch.finfo().min. For torch.float32 the min value is -3.4028234663852886e+38 which is smaller than the bfloat16 minimum of -3.3895313892515355e+38. So the problem is that torch.finfo(torch.float32).min = -3.4028234663852886e+38 gets converted to -inf. When the original encoder_extended_attention_mask is 1, then encoder_extended_attention_mask becomes (1.0 - 1.0 ) * -inf which becomes NaN (via IEEE rule Inf * 0.0 = NaN).

This PR ensures torch.finfo(torch.bfloat16).min = -3.3895313892515355e+38 and not -inf. Then the results would not have Nans.

The following lines checks for XLA_USE_BF16 or XLA_DOWNCAST_BF16 environment variable and sets the dtype accordingly:

if is_torch_tpu_available():
    if os.environ.get("XLA_USE_BF16") == 1:
         return torch.bfloat16
    if os.environ.get("XLA_DOWNCAST_BF16") == 1:
         if t.dtype == torch.float:
             return torch.bfloat16
         if t.dtype == torch.double:
             return torch.float32

Referencing related issues: aws-neuron/aws-neuron-sdk#593 and pytorch/xla#4152

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@sgugger